'''
双移动平均策略 Double Moving Average

顾名思义，双移动平均策略就是使用两条均线来判断股价未来的走势。
两条均线中，一条是长均线（如10日均线），另一条是短均线（如5日均线）。

这种策略基于一种假设：股价的动量回朝着短期均线的方向移动。当短期均线超过长期移动均线时，
动量向上，此时股价可能会上涨。然而，如果短期均线的移动方向相反，则股价可能下跌。
'''

import pandas as pd
import numpy as np
from chapter_2_get import gen_stock_data_table


# 双移动平均策略的封装
def gen_dma_strategy_data(symbol, start_date, end_date):
    stock_data = gen_stock_data_table(symbol, start_date, end_date)
    stock_002419 = stock_data['stock']
    # 新建一个数据表 strategy
    # 序号和原始数据保持一致
    strategy = pd.DataFrame(index=stock_002419.index)
    # signal，用来存储交易信号 卖出：0， 买入：1
    strategy['price'] = stock_002419['close']
    strategy['signal'] = 0

    # 将5日均线保存到avg_5，10日均线保存到avg_10
    # dataframe.rolling()函数提供滚动窗口计算的功能
    strategy['avg_5'] = stock_002419['close'].rolling(5).mean()
    strategy['avg_10'] = stock_002419['close'].rolling(10).mean()

    # 当5日均价大于10日均价时，标记为1---买入
    # 反之，标记为0---卖出
    strategy['signal'] = np.where(strategy['avg_5'] > strategy['avg_10'], 1, 0)

    # 根据信号的变化下单，当信号为0-->1时买入
    # 当信号从1-->0时卖出
    # 交易信号不变时不下单
    strategy['order'] = strategy['signal'].diff()
    order_cn_map = {
        -1.0: '卖出',
        0.0: '持有',
        1.0: '买入',
    }
    strategy['order_cn'] = strategy['order'].map(order_cn_map)
    # strategy = strategy.fillna('--')
    return strategy
